import tensorflow as tf
from tensorflow.contrib import slim
from tensorflow.contrib.tensorboard.plugins import projector

import numpy as np

import cv2

from enum import IntEnum

#交并比
def IoU(box, boxes):

    box_area = (box[2] - box[0] + 1) * (box[3] - box[1] + 1)
    area = (boxes[:, 2] - boxes[:, 0] + 1) * (boxes[:, 3] - boxes[:, 1] + 1)
    xx1 = np.maximum(box[0], boxes[:, 0])
    yy1 = np.maximum(box[1], boxes[:, 1])
    xx2 = np.minimum(box[2], boxes[:, 2])
    yy2 = np.minimum(box[3], boxes[:, 3])

    # compute the width and height of the bounding box
    w = np.maximum(0, xx2 - xx1 + 1)
    h = np.maximum(0, yy2 - yy1 + 1)

    inter = w * h
    ovr = inter / (box_area + area - inter)
    return ovr


def bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


class lineType(IntEnum):
    fileName=1,
    boxesNum=2,
    boxesCoord=3,
    

# def writeNetDataPR(tfFileName,im_size):
#     with open('./wider_face_split/wider_face_train_bbx_gt.txt','r') as f:
#         trainImagesDataDesc = f.readlines()
    
#     lines = trainImagesDataDesc.readLines()

#     print('Writing:' + tfFileName)
#     writer = tf.python_io.TFRecordWriter(tfFileName)  

#     currentLineType = lineType.fileName 
    
 
#     for line in lines:
#         if currentLineType == lineType.fileName :
#             fileName = line
#             im = cv2.imread('./WIDER_train/images/' +fileName)
#             im = cv2.resize(im,(im_size,im_size))
#             im = im.astype('uint8')
#             image_raw = im.tostring()
#             currentLineType = lineType.boxesNum
#         elif currentLineType ==lineType.boxesNum:
#             boxNum = int(line)
#             boxes = []
#             boxCounter = 0
#             currentLineType = lineType.boxesCoord
#         elif currentLineType ==lineType.boxesNum:
#             word = line.split(' ')
#             x1 = float(word[0])
#             y1 = float(word[0])
#             w = float(word[0])
#             h = float(word[0])
#             xp1 = x1 / w
#             yp1 = y1 / h
#             xp2 = (x1 +w) / w
#             yp2 = (y1 +w) / w
#             boxes.append([xp1,yp1,xp2,yp2])
#             boxCounter += 1
#             if boxCounter >= boxNum:
#                 currentLineType = lineType.fileName
#                 #生成IoU
#                 boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
#                 for box in boxes:
#                     iou = IoU(box, boxes)
#                     classification = np.array([1], dtype='float32')

                    


def writeNetData(labelFile, tfFileName,im_size):
    #读取训练集labeltxt
    
    #负样本，正样本，部分样本，关键点样本比例=3:1:1:2

    with open(labelFile,'r') as f:
        trainImagesDataDesc = f.readlines()

    print('Writing:' + tfFileName)
    writer = tf.python_io.TFRecordWriter(tfFileName)    

    examples = []

    for line in trainImagesDataDesc:

        descs = line.split()
        filePath = descs[0]

        # print(line)
        image_file_path = './' +filePath
        image = cv2.imread(image_file_path)
        h,w,ch = image.shape
        image = image.astype('uint8')
        # print(image.shape)
        

        # x1, x2, y1, y2 -> x1, y1, x2, y2

        bb = [float(descs[3]) / h,float(descs[1]) / w, float(descs[4]) / h, float(descs[2]) / w]


        # im = image.copy()
        # # x1, y1, w, h
        # bb = [float(descs[1]) / w,
        #         float(descs[3]) / h,
        #         float(descs[2]) / w  -float(descs[1]) / w ,
        #         float(descs[4]) / h - float(descs[3]) / h]
        # # 1-x1
        # # 2-y1
        # # 3-x2
        # # 4-y2
        # cv2.rectangle(im,
        #             (int(descs[1]),int(descs[3])),
        #             (int(descs[2]), int(descs[4])), 
        #             (0,255,0),3)  


        # cv2.rectangle(im,
        #             (int(float(descs[5])),int(float(descs[6]))),
        #             (int(float(descs[5]) + 2), int(float(descs[6])) + 2), 
        #             (0,255,0),3)     
        # cv2.rectangle(im,
        #             (int(float(descs[7])),int(float(descs[8]))),
        #             (int(float(descs[7]) + 2), int(float(descs[8])) + 2), 
        #             (0,255,0),3)                                        
        # cv2.rectangle(im,
        #             (int(float(descs[9])),int(float(descs[10]))),
        #             (int(float(descs[9]) + 2), int(float(descs[10])) + 2), 
        #             (0,255,0),3)     
        # cv2.rectangle(im,
        #             (int(float(descs[11])),int(float(descs[12]))),
        #             (int(float(descs[11]) + 2), int(float(descs[12])) + 2), 
        #             (0,255,0),3)     
        # cv2.rectangle(im,
        #             (int(float(descs[13])),int(float(descs[14]))),
        #             (int(float(descs[13])) + 2, int(float(descs[14])) + 2), 
        #             (0,255,0),3)     
        
        # cv2.imwrite('./rectangle/' + filePath,im)
                        

        # 5 对
        landmark = [float(descs[6])/ h,
                    float(descs[5])/ w,
                    float(descs[8])/ h,
                    float(descs[7])/ w,
                    float(descs[10])/ h,
                    float(descs[9])/ w,
                    float(descs[12])/ h,
                    float(descs[11])/ w,
                    float(descs[14])/ h,
                    float(descs[13])/ w]
        
        

        
        

        bb = np.array(bb,dtype='float32')
        bb_raw = bb.tostring()
        
        landmark = np.array(landmark,dtype='float32')
        landmark_raw = landmark.tostring()
        

        # PRO都用
        # 正100%样本0
        _bb = bb.copy()
        _bb[0] = 0.0
        _bb[1] = 0.0
        _bb[2] = 1.0
        _bb[3] = 1.0    
        im_100 = image[int(bb[0] * h):int(bb[2] * h), int(bb[1] * w):int(bb[3] * w),:]
        # cv2.imwrite("./bbr_crop/" + filePath, im_100)
        im_100 = cv2.resize(im_100,(im_size,im_size))
        _landmark = landmark.copy()
        _landmark[[0,2,4,6,8]] = (_landmark[[0,2,4,6,8]] - bb[0])/ (bb[2] -bb[0])
        _landmark[[1,3,5,7,9]] = (_landmark[[1,3,5,7,9]] - bb[1]) / (bb[3] -bb[1])
        # 第一位代表没有脸，第二位代表有脸
        classification = np.array([0,1], dtype='float32')
        cls_raw = classification.tostring()
        example = tf.train.Example(features = tf.train.Features(feature={
            'cls_raw':bytes_feature(cls_raw),
            'bb_raw':bytes_feature(_bb.tostring()),
            'landmark_raw':bytes_feature(_landmark.tostring()),
            'image_raw':bytes_feature(im_100.tostring())
            }))
        examples.append(example)
        


        if h != im_size or w != im_size:
            im = cv2.resize(image, (im_size,im_size))

        image_raw = im.tostring()
        #Pnet
        # 50%正样本
        if im_size == 12:
            classification = np.array([0,1], dtype='float32')
            cls_raw = classification.tostring()
            example = tf.train.Example(features = tf.train.Features(feature={
                'cls_raw':bytes_feature(cls_raw),
                'bb_raw':bytes_feature(bb.tostring()),
                'landmark_raw':bytes_feature(landmark.tostring()),
                'image_raw':bytes_feature(image_raw)
                }))
            examples.append(example)

       
        # 再做2个负样本
        classification = np.array([1,0], dtype='float32')
        cls_raw = classification.tostring()

        bb_ = bb.copy()
        landmark_ = landmark.copy()

        bb_[0] = 0.0
        bb_[1] = 0.0
        bb_[2] = 0.0
        bb_[3] = 0.0
        landmark_[0] = 0.0
        landmark_[1] = 0.0
        landmark_[2] = 0.0
        landmark_[3] = 0.0
        landmark_[4] = 0.0
        landmark_[5] = 0.0
        landmark_[6] = 0.0
        landmark_[7] = 0.0
        #左上角
        x1_ = 0
        y1_ = 0
        x2_ = int(bb[0] * h)
        y2_ = int(bb[1] * w)
        im_crop = image[x1_:x2_, y1_:y2_]
        # print(im_crop.shape)
        im = cv2.resize(im_crop,(im_size,im_size))
        classification = np.array([1,0], dtype='float32')
        example = tf.train.Example(features = tf.train.Features(feature={
            'cls_raw':bytes_feature(classification.tostring()),
            'bb_raw':bytes_feature(bb_.tostring()),
            'landmark_raw':bytes_feature(landmark_.tostring()),
            'image_raw':bytes_feature(im.tostring())
            }))
        examples.append(example)
    
        # if im_size == 12:
        #右下角
        x1_ = int(bb[2] * h)
        y1_ = int(bb[3] * w)
        x2_ = h
        y2_ = w
        im_crop = image[x1_:x2_, y1_:y2_]
        # print(x1_)
        # print(y1_)
        # print(x2_)
        # print(y2_)
        # print(bb)
        im = cv2.resize(im_crop,(im_size,im_size))
        classification = np.array([1,0], dtype='float32')
        example = tf.train.Example(features = tf.train.Features(feature={
            'cls_raw':bytes_feature(classification.tostring()),
            'bb_raw':bytes_feature(bb_.tostring()),
            'landmark_raw':bytes_feature(landmark_.tostring()),
            'image_raw':bytes_feature(im.tostring())
            }))
        examples.append(example)

            # # print(np.zeros(4))
            # #右上角
            # x1_ = int(bb[2] * w)
            # y1_ = 0
            # x2_ = w
            # y2_ = h
            # im_crop = image[x1_:x2_, y1_:y2_]
            # im = cv2.resize(im_crop,(im_size,im_size))
            # classification = np.array([1,0], dtype='float32')
            # example = tf.train.Example(features = tf.train.Features(feature={
            #     'cls_raw':bytes_feature(classification.tostring()),
            #     'bb_raw':bytes_feature(bb_.tostring()),
            #     'landmark_raw':bytes_feature(landmark_.tostring()),
            #     'image_raw':bytes_feature(im.tostring())
            #     }))
            # examples.append(example)
        

        
    print(len(examples))
    for example in examples:
        writer.write(example.SerializeToString())
    writer.close()   

def main():

    # writeNetData('trainImageList.txt','pnet_data.tfrecords',12)

    writeNetData('trainImageList.txt','pnet_data.tfrecords',12)
    writeNetData('trainImageList.txt','rnet_data.tfrecords',24)
    writeNetData('trainImageList.txt','onet_data.tfrecords',48)

    



    
        
if __name__ == '__main__':
    main()